import torch
import torch.nn as nn
import numpy as np
from yolov5.models.yolo import Model
from yolov5.utils.general import check_img_size
from yolov5.utils.torch_utils import select_device

class ChannelPruner:
    def __init__(self, model, device='cuda'):
        self.model = model.to(device)
        self.device = device
        self.pruning_ratios = {}
        self.masks = {}

    def calculate_importance(self, layer):
        """计算卷积层通道重要性（L1范数）"""
        if isinstance(layer, nn.Conv2d):
            return torch.norm(layer.weight.data, p=1, dim=(1, 2, 3))
        return None

    def prune_layer(self, layer_name, pruning_ratio):
        """剪枝指定层"""
        # 获取目标层
        target_layer = dict(self.model.named_modules())[layer_name]
        if not isinstance(target_layer, nn.Conv2d):
            print(f"跳过非卷积层: {layer_name}")
            return

        # 计算通道重要性
        importance = self.calculate_importance(target_layer)
        if importance is None:
            return

        # 确定要保留的通道
        num_channels = target_layer.out_channels
        num_prune = int(num_channels * pruning_ratio)
        if num_prune <= 0:
            return

        # 获取要保留的通道索引
        _, indices = torch.sort(importance, descending=True)
        keep_indices = indices[:num_channels - num_prune]
        prune_indices = indices[num_channels - num_prune:]

        # 创建掩码
        mask = torch.ones(num_channels, device=self.device)
        mask[prune_indices] = 0
        self.masks[layer_name] = mask

        # 剪枝卷积层权重
        target_layer.weight.data = target_layer.weight.data[keep_indices]
        if target_layer.bias is not None:
            target_layer.bias.data = target_layer.bias.data[keep_indices]
        target_layer.out_channels = num_channels - num_prune

        # 更新下一层的输入通道
        self.update_next_layer(layer_name, keep_indices)

    def update_next_layer(self, current_layer_name, keep_indices):
        """更新下一层的输入通道以匹配当前层的输出通道"""
        modules = list(self.model.named_modules())
        current_idx = None
        for i, (name, _) in enumerate(modules):
            if name == current_layer_name:
                current_idx = i
                break

        if current_idx is None or current_idx + 1 >= len(modules):
            return

        next_name, next_module = modules[current_idx + 1]
        if isinstance(next_module, nn.Conv2d):
            # 更新卷积层权重
            next_module.weight.data = next_module.weight.data[:, keep_indices]
            next_module.in_channels = len(keep_indices)
        elif isinstance(next_module, nn.BatchNorm2d):
            # 更新批归一化层
            next_module.weight.data = next_module.weight.data[keep_indices]
            next_module.bias.data = next_module.bias.data[keep_indices]
            next_module.running_mean = next_module.running_mean[keep_indices]
            next_module.running_var = next_module.running_var[keep_indices]
            next_module.num_features = len(keep_indices)

    def prune_model(self, pruning_ratios):
        """剪枝整个模型"""
        self.pruning_ratios = pruning_ratios
        for layer_name, ratio in pruning_ratios.items():
            self.prune_layer(layer_name, ratio)
        return self.model

    def save_pruned_model(self, path):
        """保存剪枝后的模型"""
        torch.save({
            'model': self.model,
            'masks': self.masks,
            'pruning_ratios': self.pruning_ratios
        }, path)
        print(f"剪枝后的模型已保存至 {path}")


def prune_yolov5(model_path, output_path, pruning_ratios, img_size=640, device=''):
    """剪枝YOLOv5模型的主函数"""
    device = select_device(device)
    img_size = check_img_size(img_size)

    # 加载模型
    model = Model('models/yolov5s.yaml').to(device)
    weights = torch.load(model_path, map_location=device)['model'].float()
    model.load_state_dict(weights.state_dict())
    model.eval()

    # 创建剪枝器并剪枝模型
    pruner = ChannelPruner(model, device)
    pruned_model = pruner.prune_model(pruning_ratios)

    # 保存剪枝后的模型
    pruner.save_pruned_model(output_path)

    return pruned_model

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, required=True, help='输入模型权重路径')
    parser.add_argument('--output', type=str, required=True, help='剪枝后模型保存路径')
    parser.add_argument('--ratios', type=dict, default={'model.0.conv': 0.2, 'model.1.conv': 0.2}, help='剪枝比例')
    parser.add_argument('--img-size', type=int, default=640, help='图像大小')
    parser.add_argument('--device', default='', help='设备')
    args = parser.parse_args()

    prune_yolov5(args.weights, args.output, args.ratios, args.img_size, args.device)